from __future__ import print_function

import hashlib
import os
import platform
import time
from ctypes import CFUNCTYPE, c_double, c_int64
from distutils.sysconfig import get_config_var
from typing import Tuple, List

import llvmlite.binding as llvm

# All these initializations are required for code generation!
from llvmlite import ir

llvm.initialize()
llvm.initialize_native_target()
llvm.initialize_native_asmprinter()  # yes, even this one

suffix = get_config_var('EXT_SUFFIX')

ROOT_PATH = os.path.dirname(__file__)


def get_dll_abso_path(dll_name: str) -> str:
    for root, dirs, files in os.walk(os.path.join(ROOT_PATH, 'MelodieAllegro')):
        for f in files:
            if f.startswith(dll_name + ".") and (f.endswith(suffix)):
                print(os.path.join(root, f))
                return os.path.join(root, f)
    raise FileNotFoundError


def create_execution_engine():
    """
    Create an ExecutionEngine suitable for JIT code generation on
    the host CPU.  The engine is reusable for an arbitrary number of
    modules.
    """
    # Create a target machine representing the host
    target = llvm.Target.from_default_triple()
    # target_machine = target.create_target_machine()
    # And an execution engine with an empty backing module
    backing_mod = llvm.parse_assembly("")
    llvm.load_library_permanently(get_dll_abso_path('runtime'))
    llvm.load_library_permanently(get_dll_abso_path('random'))
    # engine = llvm.create_mcjit_compiler(backing_mod, target_machine)
    # return engine


def get_from_cache(llvm_ir: str):
    """
    在当前目录下创建缓存文件夹。

    :param llvm_ir:
    :return:
    """
    hl = hashlib.md5()

    hl.update(llvm_ir.encode(encoding='utf-8'))
    md5_hash = hl.hexdigest()
    cache_folder = '_melodie_allegro_cache'
    cache_file = os.path.join(cache_folder, md5_hash + ".o")
    if not os.path.exists(cache_folder):
        os.mkdir(cache_folder)

    if os.path.exists(cache_file):
        with open(cache_file, 'rb') as f:
            return f.read()
    else:
        return None


def save_to_cache(llvm_ir: str, bts: bytes):
    hl = hashlib.md5()

    hl.update(llvm_ir.encode(encoding='utf-8'))
    md5_hash = hl.hexdigest()
    cache_folder = '_melodie_allegro_cache'
    cache_file = os.path.join(cache_folder, md5_hash + ".o")
    if not os.path.exists(cache_folder):
        os.mkdir(cache_folder)
    with open("test.ir", 'w') as f:
        f.write(llvm_ir)

    with open(cache_file, 'wb') as f:
        return f.write(bts)


def compile_ir(llvm_ir):
    """
    Compile the LLVM IR string with the given engine.
    The compiled module object is returned.
    """
    # Create a LLVM module object from the IR
    t0 = time.time()
    mod = llvm.parse_assembly(llvm_ir)
    mod.verify()

    llvm.load_library_permanently(get_dll_abso_path('runtime'))
    llvm.load_library_permanently(get_dll_abso_path('random'))

    target = llvm.Target.from_default_triple()
    target_machine = target.create_target_machine()
    engine = llvm.create_mcjit_compiler(mod, target_machine)

    print('before-engine-time', time.time() - t0)

    def on_compiled(_mod, obj_bytes: bytes):
        save_to_cache(llvm_ir, obj_bytes)

    def get_buffer(_mod):
        return get_from_cache(llvm_ir)

    # engine.set_object_cache(on_compiled, get_buffer)
    t0 = time.time()
    engine.finalize_object()
    print('finalizing time', time.time() - t0)

    # Now add the module and make sure it is ready for execution
    # engine.add_module(mod)
    # engine.finalize_object()
    engine.run_static_constructors()
    return engine


# engine = create_execution_engine()
class Runner:
    def __init__(self, ir: str, func_name: str, func_args_type: Tuple, ret_type: any):
        t0 = time.time()
        self.engine = compile_ir(ir)
        t01 = time.time()
        print("Running init function!")
        init_func_ptr = self.engine.get_function_address('runtime_init')

        self.init_func = CFUNCTYPE(c_int64, )(init_func_ptr)
        print("Init function got!", init_func_ptr)
        self.init_func()

        self.func_ptr = self.engine.get_function_address(func_name)

        t1 = time.time()
        print("Compiling time", t1 - t0, 'compile_ir', t01 - t0)
        # Run the function via ctypes
        self.cfunc = CFUNCTYPE(ret_type, *func_args_type)(self.func_ptr)
        # self.cfunc(*func_args_value)
        # self.run()
        # return cfunc

    def run(self, args):
        self.cfunc(*args)


def pre_run(ir: str, func_name: str, func_args_type: Tuple, ret_type: any, func_args_value: Tuple):
    t0 = time.time()
    engine = compile_ir(ir)
    print("Running init function!")
    init_func_ptr = engine.get_function_address('runtime_init')
    # init_func = CFUNCTYPE(c_int64, *(c_int64,))(init_func_ptr)
    init_func = CFUNCTYPE(c_int64, )(init_func_ptr)
    print("Init function got!", init_func_ptr)
    init_func()

    func_ptr = engine.get_function_address(func_name)

    t1 = time.time()
    print("Compiling time", t1 - t0)
    # Run the function via ctypes
    cfunc = CFUNCTYPE(ret_type, *func_args_type)(func_ptr)
    cfunc(*func_args_value)
    return cfunc


def run_func(ir: str, func_name: str, func_args_type: Tuple, ret_type: any, func_args_value: Tuple, statistics=None):
    t0 = time.time()
    engine = compile_ir(ir)
    print("Running init function!")
    init_func_ptr = engine.get_function_address('runtime_init')
    # init_func = CFUNCTYPE(c_int64, *(c_int64,))(init_func_ptr)
    init_func = CFUNCTYPE(c_int64, )(init_func_ptr)
    print("Init function got!", init_func_ptr)
    init_func()

    func_ptr = engine.get_function_address(func_name)

    t1 = time.time()
    print("Compiling time", t1 - t0)
    # Run the function via ctypes
    cfunc = CFUNCTYPE(ret_type, *func_args_type)(func_ptr)
    assert len(func_args_value) == len(func_args_type)
    print("before run...")
    res = cfunc(*func_args_value)
    t2 = time.time()
    if isinstance(statistics, dict):
        statistics['time'] = t2 - t1
    print("fpadd(...) =", res)
    return res


from testcall import main


def a():
    m = ir.Module()
    func_ty = ir.FunctionType(ir.VoidType(), [ir.IntType(64)])
    i32_ty = ir.IntType(32)
    func = ir.Function(m, func_ty, name="printer")

    voidptr_ty = ir.IntType(8).as_pointer()

    fmt = "Hello, %s! %i times!\n\0"
    c_fmt = ir.Constant(ir.ArrayType(ir.IntType(8), len(fmt)),
                        bytearray(fmt.encode("utf8")))
    global_fmt = ir.GlobalVariable(m, c_fmt.type, name="fstr")
    global_fmt.linkage = 'internal'
    global_fmt.global_constant = True
    global_fmt.initializer = c_fmt

    arg = "World\0"
    c_str_val = ir.Constant(ir.ArrayType(ir.IntType(8), len(arg)),
                            bytearray(arg.encode("utf8")))

    printf_ty = ir.FunctionType(ir.IntType(32), [voidptr_ty], var_arg=True)
    printf = ir.Function(m, printf_ty, name="printf")

    print1ty = ir.FunctionType(ir.IntType(64), [ir.IntType(64)])
    print1 = ir.Function(m, print1ty, name="pm_runtime_init")

    builder = ir.IRBuilder(func.append_basic_block('entry'))

    c_str = builder.alloca(c_str_val.type)
    builder.store(c_str_val, c_str)

    # this val can come from anywhere
    int_val = builder.add(i32_ty(5), i32_ty(3))

    fmt_arg = builder.bitcast(global_fmt, voidptr_ty)
    builder.call(printf, [fmt_arg, c_str, int_val])
    builder.call(print1, [ir.IntType(64)(1)])
    builder.ret_void()
    return m


def run_f(ir_code):
    main(ir_code, func_name='runtime_init')


if __name__ == "__main__":
    with open(r"C:\Users\12957\Documents\Developing\Python\llvm-py-demo\tests\temp.ll") as f:
        run_f(f.read())
    # with open('a123.ll') as f:
    #     llvm_ir = f.read()
    # run_func(llvm_ir, "nested_while_loop", (c_double,), c_double, (100.0,))
